from typing import List


# matrix for mix columns
colM =  [ 
            2, 3, 1, 1,
            1, 2, 3, 1,
            1, 1, 2, 3,
            3, 1, 1, 2
        ]

# matrix for Inverse mix columns
deColM = [ 
            0xe, 0xb, 0xd, 0x9,
            0x9, 0xe, 0xb, 0xd,
            0xd, 0x9, 0xe, 0xb,
            0xb, 0xd, 0x9, 0xe
         ]   

# Mix columns
class MC:
    def __init__(self) -> None:
        pass
    
    # Mix columns
    def encrypt(self, nums: List[List[int]]) -> List[List[int]]:
        n = [[0,0,0,0] for _ in range(4)]
        for i in range(len(nums)):
            for j in range(len(nums)):
                mul = ([self.GFMul(colM[i*4+k], nums[k][j]) for k in range(len(nums))])
                n[i][j] = mul[0] ^ mul[1] ^ mul[2] ^ mul[3]
        return n

    # Inverse Mix columns
    def decrypt(self, nums: List[List[int]]) -> List[List[int]]:
        n = [[0,0,0,0] for _ in range(4)]
        for i in range(len(nums)):
            for j in range(len(nums)):
                mul = ([self.GFMul(deColM[i*4+k], nums[k][j]) for k in range(len(nums))])
                n[i][j] = mul[0] ^ mul[1] ^ mul[2] ^ mul[3]
        return n

    # calculate x * f(x) in finitefield
    def GFMULx(self, f: int) -> int:
        res = f << 1
        a7 = res & 0x100
        if a7 != 0:
            res &= 0xff
            res ^= 0x1b

        return res

    # calculate f1(x) * f2(x) in finitefield
    def GFMul(self,f1: int,f2: int) -> int:
        f1 &= 0x0f
        f2 &= 0xff
        sum = f2
        res = 0

        while f1 != 0:
            if f1 & 1:   res ^= sum
            f1 >>= 1
            sum = self.GFMULx(sum)
        return res

# for debug
# n = [[110, 35, 54, 30], [251, 212, 204, 227], [83, 89, 79, 82], [3, 29, 60, 192]]
# print(n)
# n = MC().encrypt(n)
# print(n)
# n = MC().decrypt(n)
# print(n)
# print(MC().GFMul(0x59,0x15))